# prepare_data.py
import argparse
import os
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

def main():
    parser = argparse.ArgumentParser(description="下载并保存 Hugging Face 模型和数据集到本地")
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        required=True,
        help="要下载的 HuggingFace 模型名称或路径。"
    )
    parser.add_argument(
        "--model_output_dir",
        type=str,
        default="./data/model",
        help="保存下载模型的本地目录。"
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        required=True,
        help="要下载的 HuggingFace 数据集名称。"
    )
    parser.add_argument(
        "--dataset_config_name",
        type=str,
        default=None,
        help="HuggingFace 数据集配置名称 (可选)。"
    )
    parser.add_argument(
        "--dataset_output_dir",
        type=str,
        default="./data/dataset",
        help="保存下载数据集的本地目录。"
    )
    parser.add_argument(
        "--cache_dir",
        type=str,
        default=None,
        help="Hugging Face Hub 下载缓存目录 (可选)。"
    )

    args = parser.parse_args()

    print(f"准备参数: {args}")

    # 创建输出目录 (如果不存在)
    os.makedirs(args.model_output_dir, exist_ok=True)
    os.makedirs(args.dataset_output_dir, exist_ok=True)

    # 1. 下载并保存模型和 Tokenizer
    print(f"正在从 '{args.model_name_or_path}' 下载模型和 Tokenizer...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
        model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
        
        print(f"正在将模型和 Tokenizer 保存到 '{args.model_output_dir}'...")
        tokenizer.save_pretrained(args.model_output_dir)
        model.save_pretrained(args.model_output_dir)
        print("模型和 Tokenizer 保存成功。")
    except Exception as e:
        print(f"下载或保存模型 '{args.model_name_or_path}' 失败: {e}")
        # 可以在这里决定是否继续下载数据集，或者直接退出

    # 2. 下载并保存数据集
    print(f"正在下载数据集 '{args.dataset_name}' (配置: '{args.dataset_config_name}')...")
    try:
        dataset = load_dataset(
            args.dataset_name, 
            name=args.dataset_config_name, 
            cache_dir=args.cache_dir
        )
        print(f"正在将数据集保存到 '{args.dataset_output_dir}'...")
        # datasets 库的 save_to_disk 方法可以将数据集保存到磁盘
        # 它会为每个 split (train, validation, test 等) 创建子目录
        dataset.save_to_disk(args.dataset_output_dir)
        print("数据集保存成功。")
    except Exception as e:
        print(f"下载或保存数据集 '{args.dataset_name}' 失败: {e}")

    print("数据准备完成。")

if __name__ == "__main__":
    main()